# -*- coding: utf-8 -*-#

# -------------------------------------------------------------------------------
# Name:         dataset
# Description:
# Author:       zx
# Date:         2020/9/28
# -------------------------------------------------------------------------------
import tensorflow as tf
import numpy as np
import cv2
import main


# 函数的功能时将filename对应的图片文件读进来，并缩放到统一的大小
def _parse_function(filename, label_file):
    image_string = tf.read_file(filename)
    image_decoded = tf.image.decode_jpeg(image_string, channels=3)
    image_resized = tf.image.resize_images(image_decoded, [28, 28])
    return {"image_resized": image_resized.shape[0], "label_file": label_file}


def read_object_caller(filename, label_file):
    # 使用tf.py_func调用一个普通python函数来读取一个物体的12张图片路径
    # 注意返回值的类型是[tf.string, label.dtype]。
    return tf.py_func(read_object_list, [filename, label_file], [tf.tuple])


def read_object_list(filename, label):
    # print(dict(main.main().__getitem__(1)))
    # return dict(main.main().__getitem__(1))
    image = cv2.imread(str(filename, encoding="utf-8"))
    return {"image": image, "label": label}


# 图片文件的列表
filenames = tf.constant([r"datasets/total_text/train_images/img11.jpg",
                         r"datasets/total_text/train_images/img12.jpg"])
# label[i]就是图片filenames[i]的label
label_files = tf.constant([r"datasets/total_text/train_gts/img11.jpg.txt",
                           r"datasets/total_text/train_gts/img12.jpg.txt"])

# 此时dataset中的一个元素是(filename, label)
dataset = tf.data.Dataset.from_tensor_slices((filenames, label_files))

# 此时dataset中的一个元素是(image_resized, label)
# dataset = dataset.map(_parse_function)
dataset = dataset.map(read_object_caller)

# 此时dataset中的一个元素是(image_resized_batch, label_batch)
dataset = dataset.shuffle(2).batch(1).repeat(1)

iterator = dataset.make_one_shot_iterator()
one_element = iterator.get_next()
with tf.Session() as sess:
    try:
        while True:
            print(sess.run(one_element))
    except tf.errors.OutOfRangeError:
        print("end!")
